import cv2
import numpy as np
import json
import os
from PyQt5.QtWidgets import QDialog, QVBoxLayout, QHBoxLayout, QPushButton, QLabel, QTextEdit, QMessageBox
from PyQt5.QtCore import Qt, pyqtSignal
from PyQt5.QtGui import QPixmap, QImage, QPainter, QPen, QColor

class CalibrationDialog(QDialog):
    """标定比例尺对话框"""
    
    # 定义信号
    calibration_completed = pyqtSignal(float)  # 标定完成信号，传递比例尺系数
    
    def __init__(self, current_frame, parent=None):
        """
        初始化标定对话框
        
        Args:
            current_frame: 当前图像帧
            parent: 父窗口
        """
        super().__init__(parent)
        self.current_frame = current_frame.copy() if current_frame is not None else None
        self.clicked_points = []  # 存储点击的像素坐标
        self.scale_factor = 1.0   # 比例尺系数
        self.real_distance = 0.0  # 实际距离（米）
        
        self.init_ui()
        self.setup_image()
    
    def init_ui(self):
        """初始化用户界面"""
        self.setWindowTitle("标定比例尺")
        self.setModal(True)
        self.resize(800, 600)
        
        # 创建主布局
        main_layout = QVBoxLayout()
        
        # 创建图像显示区域
        self.image_label = QLabel()
        self.image_label.setAlignment(Qt.AlignCenter)
        self.image_label.setMinimumSize(640, 480)
        self.image_label.setStyleSheet("border: 2px solid gray;")
        main_layout.addWidget(self.image_label)
        
        # 创建说明标签
        instruction_label = QLabel("请在图像上点击两个点来测量距离，然后输入实际距离进行标定")
        instruction_label.setAlignment(Qt.AlignCenter)
        main_layout.addWidget(instruction_label)
        
        # 创建按钮布局
        button_layout = QHBoxLayout()
        
        # 清除按钮
        self.clear_button = QPushButton("清除点")
        self.clear_button.clicked.connect(self.clear_points)
        button_layout.addWidget(self.clear_button)
        
        # 输入实际距离的标签和输入框
        distance_label = QLabel("实际距离(米):")
        button_layout.addWidget(distance_label)
        
        from PyQt5.QtWidgets import QLineEdit
        self.distance_input = QLineEdit()
        self.distance_input.setPlaceholderText("输入两点间的实际距离")
        self.distance_input.setMaximumWidth(150)
        button_layout.addWidget(self.distance_input)
        
        # 标定按钮
        self.calibrate_button = QPushButton("计算比例尺")
        self.calibrate_button.clicked.connect(self.calculate_scale)
        self.calibrate_button.setEnabled(False)
        button_layout.addWidget(self.calibrate_button)
        
        # 确认按钮
        self.confirm_button = QPushButton("确认")
        self.confirm_button.clicked.connect(self.confirm_calibration)
        self.confirm_button.setEnabled(False)
        button_layout.addWidget(self.confirm_button)
        
        # 取消按钮
        self.cancel_button = QPushButton("取消")
        self.cancel_button.clicked.connect(self.reject)
        button_layout.addWidget(self.cancel_button)
        
        main_layout.addLayout(button_layout)
        
        # 创建坐标显示区域
        self.coordinates_text = QTextEdit()
        self.coordinates_text.setMaximumHeight(100)
        self.coordinates_text.setReadOnly(True)
        self.coordinates_text.setPlaceholderText("点击的坐标将显示在这里...")
        main_layout.addWidget(self.coordinates_text)
        
        self.setLayout(main_layout)
        
        # 连接鼠标点击事件
        self.image_label.mousePressEvent = self.on_image_clicked
    
    def setup_image(self):
        """设置图像显示"""
        if self.current_frame is None:
            self.image_label.setText("无图像数据")
            return
        
        try:
            # 将OpenCV图像转换为Qt图像
            height, width, channel = self.current_frame.shape
            bytes_per_line = 3 * width
            q_image = QImage(self.current_frame.data, width, height, bytes_per_line, QImage.Format_RGB888).rgbSwapped()
            
            # 创建QPixmap并缩放以适应标签
            pixmap = QPixmap.fromImage(q_image)
            
            # 计算缩放比例以保持宽高比
            label_size = self.image_label.size()
            scaled_pixmap = pixmap.scaled(label_size, Qt.KeepAspectRatio, Qt.SmoothTransformation)
            
            self.image_label.setPixmap(scaled_pixmap)
            
            # 存储原始图像和缩放信息
            self.original_pixmap = pixmap
            self.scaled_pixmap = scaled_pixmap
            # 避免除零错误
            if scaled_pixmap.width() > 0 and scaled_pixmap.height() > 0:
                self.scale_ratio_x = pixmap.width() / scaled_pixmap.width()
                self.scale_ratio_y = pixmap.height() / scaled_pixmap.height()
            else:
                self.scale_ratio_x = 1.0
                self.scale_ratio_y = 1.0
            
        except Exception as e:
            self.image_label.setText(f"图像显示错误: {str(e)}")
    
    def on_image_clicked(self, event):
        """图像点击事件处理"""
        if self.current_frame is None or not hasattr(self, 'scaled_pixmap'):
            return
        
        # 获取点击位置
        pos = event.pos()
        
        # 计算图像在标签中的偏移量（居中显示）
        label_size = self.image_label.size()
        pixmap_size = self.scaled_pixmap.size()
        
        offset_x = (label_size.width() - pixmap_size.width()) // 2
        offset_y = (label_size.height() - pixmap_size.height()) // 2
        
        # 计算相对于图像的坐标
        image_x = pos.x() - offset_x
        image_y = pos.y() - offset_y
        
        # 检查点击是否在图像范围内
        if (0 <= image_x < pixmap_size.width() and 
            0 <= image_y < pixmap_size.height()):
            
            # 转换为原始图像坐标
            original_x = int(image_x * self.scale_ratio_x)
            original_y = int(image_y * self.scale_ratio_y)
            
            # 添加点击点
            self.clicked_points.append((original_x, original_y))
            
            # 更新显示
            self.update_coordinates_display()
            self.update_image_with_points()
            
            # 检查是否可以启用标定按钮
            if len(self.clicked_points) >= 2:
                self.calibrate_button.setEnabled(True)
    
    def update_coordinates_display(self):
        """更新坐标显示"""
        text = "点击的坐标:\n"
        for i, (x, y) in enumerate(self.clicked_points):
            text += f"点{i+1}: ({x}, {y})\n"
        
        if len(self.clicked_points) >= 2:
            # 计算像素距离
            p1 = self.clicked_points[0]
            p2 = self.clicked_points[1]
            pixel_distance = np.sqrt((p2[0] - p1[0])**2 + (p2[1] - p1[1])**2)
            text += f"像素距离: {pixel_distance:.2f} 像素\n"
        
        self.coordinates_text.setText(text)
    
    def update_image_with_points(self):
        """在图像上绘制点击的点"""
        if self.current_frame is None or len(self.clicked_points) == 0:
            return
        
        try:
            # 创建图像副本
            display_image = self.current_frame.copy()
            
            # 绘制所有点击的点
            for i, (x, y) in enumerate(self.clicked_points):
                # 绘制圆点
                cv2.circle(display_image, (x, y), 5, (0, 255, 0), -1)
                # 绘制序号
                cv2.putText(display_image, str(i+1), (x+10, y-10), 
                           cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2)
            
            # 如果有两个点，绘制连线
            if len(self.clicked_points) >= 2:
                p1 = self.clicked_points[0]
                p2 = self.clicked_points[1]
                cv2.line(display_image, p1, p2, (0, 0, 255), 2)
            
            # 更新显示
            height, width, channel = display_image.shape
            bytes_per_line = 3 * width
            q_image = QImage(display_image.data, width, height, bytes_per_line, QImage.Format_RGB888).rgbSwapped()
            
            pixmap = QPixmap.fromImage(q_image)
            label_size = self.image_label.size()
            scaled_pixmap = pixmap.scaled(label_size, Qt.KeepAspectRatio, Qt.SmoothTransformation)
            
            self.image_label.setPixmap(scaled_pixmap)
            
            # 更新缩放信息
            self.original_pixmap = pixmap
            self.scaled_pixmap = scaled_pixmap
            if scaled_pixmap.width() > 0 and scaled_pixmap.height() > 0:
                self.scale_ratio_x = pixmap.width() / scaled_pixmap.width()
                self.scale_ratio_y = pixmap.height() / scaled_pixmap.height()
            
        except Exception as e:
            print(f"更新图像显示错误: {e}")
    
    def clear_points(self):
        """清除所有点击的点"""
        self.clicked_points.clear()
        self.update_coordinates_display()
        self.setup_image()  # 重新显示原始图像
        self.calibrate_button.setEnabled(False)
        self.confirm_button.setEnabled(False)
    
    def calculate_scale(self):
        """计算比例尺系数"""
        if len(self.clicked_points) < 2:
            QMessageBox.warning(self, "警告", "请至少点击两个点")
            return
        
        # 获取实际距离输入
        try:
            self.real_distance = float(self.distance_input.text())
            if self.real_distance <= 0:
                QMessageBox.warning(self, "警告", "请输入有效的实际距离")
                return
        except ValueError:
            QMessageBox.warning(self, "警告", "请输入有效的数字")
            return
        
        # 计算像素距离
        p1 = self.clicked_points[0]
        p2 = self.clicked_points[1]
        pixel_distance = np.sqrt((p2[0] - p1[0])**2 + (p2[1] - p1[1])**2)
        
        # 计算比例尺系数（米/像素）
        self.scale_factor = self.real_distance / pixel_distance
        
        # 显示结果
        result_text = f"标定结果:\n"
        result_text += f"像素距离: {pixel_distance:.2f} 像素\n"
        result_text += f"实际距离: {self.real_distance:.6f} 米\n"
        result_text += f"比例尺系数: {self.scale_factor:.9f} 米/像素\n"
        result_text += f"比例尺系数: {self.scale_factor*1000:.6f} 毫米/像素"
        
        self.coordinates_text.setText(result_text)
        
        # 启用确认按钮
        self.confirm_button.setEnabled(True)
        
        QMessageBox.information(self, "标定完成", 
                              f"比例尺系数: {self.scale_factor:.9f} 米/像素\n"
                              f"({self.scale_factor*1000:.6f} 毫米/像素)")
    
    def confirm_calibration(self):
        """确认标定结果"""
        if self.scale_factor > 0:
            # 保存标定信息到文件
            if self.save_calibration_to_file(self.scale_factor):
                print("标定信息已保存到文件")
            else:
                print("标定信息保存失败")
            
            self.calibration_completed.emit(self.scale_factor)
            self.accept()
        else:
            QMessageBox.warning(self, "警告", "请先计算比例尺系数")
    
    def get_scale_factor(self):
        """获取比例尺系数"""
        return self.scale_factor
    
    def save_calibration_to_file(self, scale_factor):
        """
        保存标定信息到文件
        
        Args:
            scale_factor (float): 比例尺系数
        """
        try:
            # 创建配置目录
            config_dir = "config"
            if not os.path.exists(config_dir):
                os.makedirs(config_dir)
            
            # 标定信息文件路径
            calibration_file = os.path.join(config_dir, "calibration.json")
            
            # 准备保存的数据
            calibration_data = {
                "scale_factor": scale_factor,
                "unit": "米/像素",
                "timestamp": str(np.datetime64('now')),
                "description": "比例尺标定系数"
            }
            
            # 保存到JSON文件
            with open(calibration_file, 'w', encoding='utf-8') as f:
                json.dump(calibration_data, f, ensure_ascii=False, indent=2)
            
            print(f"标定信息已保存到: {calibration_file}")
            return True
            
        except Exception as e:
            print(f"保存标定信息失败: {e}")
            return False
    
    @staticmethod
    def load_calibration_from_file():
        """
        从文件加载标定信息
        
        Returns:
            float: 比例尺系数，如果加载失败返回None
        """
        try:
            # 标定信息文件路径
            calibration_file = os.path.join("config", "calibration.json")
            
            # 检查文件是否存在
            if not os.path.exists(calibration_file):
                print("标定文件不存在")
                return None
            
            # 从JSON文件加载
            with open(calibration_file, 'r', encoding='utf-8') as f:
                calibration_data = json.load(f)
            
            # 获取比例尺系数
            scale_factor = calibration_data.get("scale_factor")
            
            if scale_factor is not None:
                print(f"已加载标定系数: {scale_factor:.9f} 米/像素")
                return scale_factor
            else:
                print("标定文件中未找到比例尺系数")
                return None
                
        except Exception as e:
            print(f"加载标定信息失败: {e}")
            return None 